import gradio as gr
from deoldify.device_id import DeviceId
from deoldify import device
from deoldify.visualize import get_image_colorizer
from PIL import Image

# 创建colorizer实例, 参数artistic=True表示使用艺术模式
colorizer = get_image_colorizer(artistic=True)

# 定义Gradio应用的函数
def colorize_image(image, render_factor):
    # 将上传的图像保存到本地
    image = Image.fromarray(image)
    image_path = "data/uploaded_image.png"
    image.save(image_path)
    
    # 使用DeOldify进行上色处理
    result_path = colorizer.plot_transformed_image(image_path, render_factor=render_factor, compare=True)
    return result_path

# 创建Gradio界面
iface = gr.Interface(
    fn=colorize_image,  # 处理函数
    inputs=[
        gr.Image(type="numpy", label="image"),  # 上传图片控件
        gr.Slider(10, 40, step=1, label="render factor")     # 渲染因子滑动条控件
    ],
    outputs=gr.Image(label="output"),   # 输出图片控件  
    title="图片上色器",  # 应用标题
    description="上传一张黑白图片并通过滑动条来调整上色效果。"  # 应用描述
)

# 运行应用
iface.launch()
